from rest_framework import viewsets
from django_celery_beat.models import PeriodicTask, CrontabSchedule
from django.db import transaction

from task.serializers import ScheduleSerializer
from libs.decorators import required_data_fields, permission_required
from libs.enums import Permission, LogModule
from task.models import ConnectionTask
from log.services import log_service


class ScheduleViewSet(viewsets.ModelViewSet):
    queryset = PeriodicTask.objects.all().select_related('crontab').order_by('-id')
    serializer_class = ScheduleSerializer
    ordering_fields = ('id',)
    log_module = LogModule.TASK.value

    @transaction.atomic
    @permission_required([Permission.TASK_MANAGE.value])
    @required_data_fields(['name', 'task', 'crontab'])
    def create(self, request, *args, **kwargs):
        crontab_data = request.data.pop('crontab')
        connection_id = request.data.pop('connection_id')
        crontab, _ = CrontabSchedule.objects.get_or_create(**crontab_data)
        request.data['crontab'] = crontab.id
        result = super().create(request, *args, **kwargs)
        if connection_id:
            ConnectionTask.objects.create(
                connection_id=connection_id,
                periodic_task_id=result.data['id']
            )
        return result

    def perform_create(self, serializer):
        super().perform_create(serializer)
        log_service.build_create_log(self.request.user, serializer.instance, self.log_module, f'创建了任务：{serializer.instance.name}')

    @permission_required([Permission.TASK_MANAGE.value])
    def update(self, request, *args, **kwargs):
        return super().update(request, *args, **kwargs)

    @permission_required([Permission.TASK_MANAGE.value])
    def partial_update(self, request, *args, **kwargs):
        crontab_data = request.data.pop('crontab', None)
        if crontab_data:
            crontab, _ = CrontabSchedule.objects.get_or_create(**crontab_data)
            request.data['crontab'] = crontab.id
        return super().partial_update(request, *args, **kwargs)

    def perform_update(self, serializer):
        super().perform_update(serializer)
        log_service.build_update_log(self.request.user, serializer.instance, self.log_module, f'更新了任务：{serializer.instance.name}')
